#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy as np
import SO3_base
import eulerangles as euler
import Quaternion as quat
import transform as tf

class SE3():
    def __init__(self):
        self.translation = np.zeros(3)
        self.rotation = np.eye(3)

    def __add__(self, other):
        self.translation += other.translation()
        self.rotation.dot(other.rotation_matrix())

    def __sub__(self, other):
        self.translation -= other.translation()
        self.rotation.dot(np.linalg.inv(other.rotation_matrix()))

    def translation(self):
        return self.translation

    def rotation_matrix(self):
        return self.rotation

    def hat(self, v):
        """
        给定向量，输入其hat矩阵
        :param
        v： 4d vector
        :return
        out: 4d skew matrix
        """
        if len(v) == 4:
            v = v[:3] / v[3]
        else:
            raise LookupError('please input proper vector')
        skv = np.roll(np.roll(np.diag(v.flatten()), 1, 1), -1, 0)
        return skv - skv.T

    def exp(self, se3):
        """
        :param
        se3: 4d se3 vector
        :return
        SE3: 4d SE3 map matrix
        """
        phi = se3[3:]  # about orientation
        rho = se3[:3]  # about translation
        theta = np.linalg.norm(phi)
        if theta == 0:
            rot = np.eye(3)
            J = np.eye(3)
        else:
            # with np.errstate(invalid='ignore'):
            alpha = phi / theta
            alpha = np.nan_to_num(alpha)
            alpha = np.expand_dims(alpha, axis=0)  # row vector
            rot = np.cos(theta) * np.eye(3) + (1 - np.cos(theta)) * alpha.T.dot(alpha) + np.sin(theta) * self.hat(
                alpha)
            J = np.sin(theta) / theta * np.eye(3) + (1 - np.sin(theta) / theta) * alpha.T.dot(alpha) + (1 - np.cos(
                theta)) / theta * self.hat(alpha)
            J = np.nan_to_num(J)
        trans = J.dot(rho)
        SE3 = np.eye(4)
        SE3[:3, :3] = rot
        SE3[:3, 3] = trans
        return SE3

    def log(self, SE3):
        '''
        :param
        SE3: 4d SE3 matrix
        :return
        se3: 4d se3 vector
        '''
        rot = SE3[:3, :3]
        trans = SE3[:3, 3]
        theta = np.arccos((np.trace(rot) - 1) / 2)
        if theta == 0:
            J = np.eye(3)
            alpha = np.zeros(3)
        else:
            lnR = theta / 2 / np.sin(theta) * (rot - rot.T)
            alpha = np.array([-lnR[1, 2], lnR[0, 2], -lnR[0, 1]])
            J = np.sin(theta) / theta * np.eye(3) + (1 - np.sin(theta) / theta) * alpha.T.dot(alpha) + (1 - np.cos(
                theta)) / theta * self.hat(alpha)
            J = np.nan_to_num(J)

        rho = J.dot(trans)
        se3 = np.zeros(6)
        se3[:3] = rho
        se3[3:] = alpha
        return se3

    def inverse(self):
        inv = np.eye(4)
        inv[:3, :3] = self.rotation.T
        inv[:3, 3] = self.translation
        return inv

    def Adjoints(self, SE3):
        Adj = np.zeros([6,6])
        Adj[:3, :3] = self.rotation
        Adj[3:, 3:] = self.rotation
        Adj[:3, 3:] = self.hat(self.translation)*self.rotation
        return Adj

    def so3(self):
        SO3 = SO3_base.SO3_base()
        return SO3.log(self.rotation)

    def se3(self):
        _se3 = self.log(self.exp())
        return _se3

    def set_translation(self, _translation):
        self.translation = _translation

    def set_euler(self, _euler):
        x, y, z = _euler
        self.rotation = euler.euler2mat(z, x, y)

    def set_quaternion(self, _quaternion):
        self.rotation = quat.quat2mat(_quaternion)

    def set_axis_angle(self, theta, axis):
        mat = tf.rotation_matrix(theta, axis)
        self.rotation = mat[:3, :3]




